import torch
import torch.nn as nn
import torch.nn.functional as f
import numpy as np
import cv2

from ..module import build_backbone, DETECTOR
from ..loss import loss_dict


@DETECTOR.register_module
class CamDetector(nn.Module):

    def __init__(self, backbone_cfg, stride, feat_size, fc_channels, losses):
        super(CamDetector, self).__init__()
        self.backbone = build_backbone(backbone_cfg)
        self.stride = stride
        self.fc_channels = fc_channels
        self.conv_last = torch.nn.Conv2d(feat_size[0], feat_size[0]*2, kernel_size=3, stride=1, padding=1)
        self.gap = nn.AvgPool2d(feat_size[1])
        self.fcs = nn.ModuleList()
        for i in range(len(fc_channels)):
            if i != 0:
                self.fcs.append(nn.Linear(fc_channels[i-1], fc_channels[i]))
            else:
                self.fcs.append(nn.Linear(feat_size[0]*2, fc_channels[i]))
        self.losses = []
        for loss in losses:
            self.losses.append(loss_dict[loss])
        self.backbone_feature = 0

    def forward(self, feature, label=None, mode='train', thres=None):
        backbone_feature = self.backbone(feature)
        self.backbone_feature = torch.nn.functional.relu(self.conv_last(backbone_feature))
        x = self.gap(self.backbone_feature)
        x = x.view([len(x), -1])
        for fc in self.fcs:
            x = fc(x)
        cls = x
        if mode == 'train':
            loss_func = self.losses[0]
            cls_loss = loss_func(cls, label)
            return cls, cls_loss
        elif mode == 'loc':
            # For see images' bkg heat map
            # heat_map = self.heatmap_predict(backbone_feature, torch.LongTensor([200]).cuda())
            heat_map = self.heatmap_predict(backbone_feature, label)
            heat_map = f.relu(torch.from_numpy(heat_map)).numpy()
            addition_data = {'heat_map': heat_map}

            # For all classes CAM
            # for i in range(self.fc_channels[-1]):
            #     i_cls = torch.zeros_like(label, dtype=torch.long, device=backbone_feature.device)
            #     i_cls[:] = i
            #     addition_data['{}_map'.format(i)] = self.heatmap_predict(backbone_feature, i_cls)

            # Feat maps
            # addition_data['featmap'] = backbone_feature.clone().detach().cpu().numpy()

            return cls, [bbox_predict(heat_map, feature.shape[-2:], threshold=t) for t in thres], addition_data
        else:
            return f.softmax(cls, dim=-1)

    def init_weights(self):
        self.backbone.init_weights()
        torch.nn.init.normal_(self.conv_last.weight, std=0.01)
        torch.nn.init.constant_(self.conv_last.bias, 0)
        for fc in self.fcs:
            torch.nn.init.normal_(fc.weight, std=0.001)
            torch.nn.init.constant_(fc.bias, 0)

    def heatmap_predict(self, backbone_feature, cls):
        weight = torch.index_select(list(self.fcs)[0].weight, 0, cls)
        return (weight.unsqueeze(-1).unsqueeze(-1) * backbone_feature).sum(dim=1).cpu().detach().numpy()


def bbox_predict(heat_map, size, threshold):
    heat_map = heat_map[0]
    heat_map = cv2.resize(heat_map, size)

    max_v = heat_map.max()
    min_v = heat_map.min()
    heat_map = (heat_map - min_v) / (max_v - min_v + 1e-8)

    binary_map = (heat_map >= threshold)

    bboxes = np.zeros(shape=(heat_map.shape[0], 4), dtype=np.uint16)
    ys, xs = np.where(binary_map == 1)
    if len(ys) == 0:
        y1, x1, y2, x2 = 0, 0, 0, 0
        print('No bbox can be found')
    else:
        y1, x1, y2, x2 = ys.min(), xs.min(), ys.max(), xs.max()
    bboxes[0, 0] = int(x1)
    bboxes[0, 1] = int(y1)
    bboxes[0, 2] = int(x2)
    bboxes[0, 3] = int(y2)

    return bboxes